import torch  # 导入 PyTorch 库，用于深度学习模型的构建和推理
from net import Model  # 导入自定义的模型类（假设基于 BERT 实现）
from transformers import BertTokenizer  # 导入 Hugging Face 的 BERT 分词器

# 定义设备信息（如果 GPU 可用，则使用 GPU；否则使用 CPU）
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)  # 打印当前使用的设备信息（GPU 或 CPU）

# 加载预训练的 BERT 分词器
token = BertTokenizer.from_pretrained(
    r"D:\SentimentBERT\model\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
)
names = ["负向评价", "正向评价"]  # 定义分类标签（负向和正向）
model = Model().to(DEVICE)  # 初始化模型，并将其移动到指定设备上

# 定义一个函数 `collate_fn`，用于对输入数据进行编码
def collate_fn(data):
    sents = []  # 创建一个空列表，用于存储输入句子
    sents.append(data)  # 将输入数据添加到列表中
    # 使用分词器对文本进行编码
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sents,  # 输入的文本列表
        truncation=True,  # 是否对超出最大长度的文本进行截断
        max_length=500,  # 设置最大长度为 500
        padding="max_length",  # 将所有文本填充到最大长度
        return_tensors="pt",  # 返回 PyTorch 张量格式的数据
        return_length=True  # 返回每个文本的实际长度
    )
    input_ids = data["input_ids"]  # 获取输入的 token ID
    attention_mask = data["attention_mask"]  # 获取 attention mask（用于忽略填充部分）
    token_type_ids = data["token_type_ids"]  # 获取 token 类型 ID（用于区分句子 A 和 B）

    return input_ids, attention_mask, token_type_ids  # 返回编码后的数据

# 定义测试函数
def test():
    # 加载训练好的模型参数
    model.load_state_dict(torch.load("params/1_bert.pth", map_location=DEVICE))
    # 开启测试模式（关闭 dropout 等训练时的行为）
    model.eval()

    while True:
        data = input("请输入测试数据（输入‘q’退出）：")  # 提示用户输入测试数据
        if data == 'q':  # 如果用户输入 'q'，则退出循环
            print("测试结束")
            break
        # 对输入数据进行编码
        input_ids, attention_mask, token_type_ids = collate_fn(data)
        # 将编码后的数据移动到指定设备上
        input_ids, attention_mask, token_type_ids = \
            input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE)

        with torch.no_grad():  # 禁用梯度计算（推理阶段不需要更新参数）
            out = model(input_ids, attention_mask, token_type_ids)  # 将数据输入模型，得到输出
            out = out.argmax(dim=1)  # 获取预测结果（取输出的最大值对应的类别）
            print("模型判定：", names[out], "\n")  # 根据预测结果打印对应的分类标签

if __name__ == '__main__':
    test()  # 调用测试函数，开始测试

    # 逻辑实现
    # 1.设备选择
    # 2.模型和分词器加载
    # 3.数据编码
    # 4.测试流程